from coin.base.book.types import TradeSide, create_market_trade_event


def _is_aggable_trade_event(t1, t2):
  if t1 is None or t2 is None:
    return False
  # qty is not part of equality
  return (t1.timestamp == t2.timestamp and t1.symbol == t2.symbol and t1.price == t2.price
          and t1.side == t2.side)


class TradeAggregator(object):
  def __init__(self, *, handler, agg_view):
    self._handler = handler
    self._agg_view = agg_view
    self._pending_trade_event = None

  def on_begin_trade(self):
    assert self._pending_trade_event is None
    if self._agg_view:
      self._pending_trade_event = None
    else:  # not self._agg_view
      pass

  def on_mid_trade(self, trade_event):
    # assert trade_event consistency
    if self._agg_view:
      if _is_aggable_trade_event(self._pending_trade_event, trade_event):
        self._pending_trade_event = create_market_trade_event(
            self._pending_trade_event.timestamp,
            self._pending_trade_event.symbol,
            self._pending_trade_event.price,
            self._pending_trade_event.qty + trade_event.qty,
            self._pending_trade_event.side,
            has_more_trade=True,
            topic_id=self._pending_trade_event.topic_id)
        # do not publish event here due to aggregation here.
        # buffer event and will be published later.
      else:
        if self._pending_trade_event is not None:
          self._handler.publish_event('trade',
                                      self._pending_trade_event.symbol,
                                      self._pending_trade_event)
        self._pending_trade_event = trade_event
    else:  # not self._agg_view
      self._handler.publish_event('trade', trade_event.symbol, trade_event)

  def on_end_trade(self):
    if self._agg_view:
      if self._pending_trade_event is not None:
        self._pending_trade_event = create_market_trade_event(
            self._pending_trade_event.timestamp,
            self._pending_trade_event.symbol,
            self._pending_trade_event.price,
            self._pending_trade_event.qty,
            self._pending_trade_event.side,
            has_more_trade=False,
            topic_id=self._pending_trade_event.topic_id)
        self._handler.publish_event('trade',
                                    self._pending_trade_event.symbol,
                                    self._pending_trade_event)
        self._pending_trade_event = None
    else:  # not self._agg_view
      pass
    assert self._pending_trade_event is None
